import json
import os
import glob
import random

# Repository list from the original code
REPO_List = [
    # "account_os_account",
    "arkui_ace_engine",
    # "build",
    # "communication_wifi",
    # "developtools_ace_ets2bundle",
    "multimedia_audio_framework",
    # "web_webview",
    # "xts_acts",
    # "kernel_linux_5.10"
]

def read_json_objects_in_file(filepath):
    """Count the number of JSON objects in a JSONL file"""
    try:
        re_list = []
        with open(filepath, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():  # Skip empty lines
                    re_list.append(json.loads(line.strip()))
        return re_list
    except FileNotFoundError:
        return f"File not found: {filepath}"
    except Exception as e:
        return f"Error reading {filepath}: {str(e)}"


print("=== RQ1&2 Repository-specific files ===")
total_repo_files = 0
for repo in REPO_List:
    OWNER = "openharmony"
    REPO_OUTPUT_FILE = f"./RQ2/{repo}/{OWNER}_{repo}_check_data.jsonl"
    json_list = read_json_objects_in_file(REPO_OUTPUT_FILE)
    random.shuffle(json_list)
    split_index = int(len(json_list) * 0.8)
    train_data = json_list[:split_index]
    test_data = json_list[split_index:]

    # 写入训练集
    TRAIN_DATA_OUTPUT_FILE = f"./RQ2/{repo}/{OWNER}_{repo}_train_data.jsonl"
    try:
        # 确保目录存在
        import os
        train_need_check_count = 0
        train_data_output_dir = os.path.dirname(TRAIN_DATA_OUTPUT_FILE)
        if not os.path.exists(train_data_output_dir):
            os.makedirs(train_data_output_dir)
        with open(TRAIN_DATA_OUTPUT_FILE, 'w', encoding='utf-8') as f:
            for item in train_data:
                if item.get('need_check') is True:
                    train_need_check_count += 1
                f.write(json.dumps(item, ensure_ascii=False) + '\n')
        print(f"已将 {len(train_data)} 条数据写入 {TRAIN_DATA_OUTPUT_FILE}, 其中 need_check 为 True 的有 {train_need_check_count} 条")
    except Exception as e:
        print(f"写入 {TRAIN_DATA_OUTPUT_FILE} 失败: {e}")

    # 写入测试集
    TEST_DATA_OUTPUT_FILE = f"./RQ2/{repo}/{OWNER}_{repo}_test_data.jsonl"
    try:
        # 确保目录存在
        import os
        test_need_check_count = 0
        test_data_output_dir = os.path.dirname(TEST_DATA_OUTPUT_FILE)
        if not os.path.exists(test_data_output_dir):
            os.makedirs(test_data_output_dir)

        with open(TEST_DATA_OUTPUT_FILE, 'w', encoding='utf-8') as f:
            for item in test_data:
                if item.get('need_check') is True:
                    test_need_check_count += 1
                f.write(json.dumps(item, ensure_ascii=False) + '\n')
        print(f"已将 {len(test_data)} 条数据写入 {TEST_DATA_OUTPUT_FILE}, 其中 need_check 为 True 的有 {test_need_check_count} 条")
    except Exception as e:
        print(f"写入 {TEST_DATA_OUTPUT_FILE} 失败: {e}")

    print(f"{REPO_OUTPUT_FILE}: {len(json_list)}")



# Count for aggregate files

